import torch
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from pandas.core.frame import DataFrame
import pandas as pd
import numpy as np
import argparse
import yaml
import matplotlib.pyplot as plt
import numba

# 用TSNE进行数据降维并展示聚类结果

parser = argparse.ArgumentParser()
parser.add_argument('--feat_file', type=str, default='./input')
parser.add_argument('--cls_file', type=str, default='./input')

args = parser.parse_args()

text_feat=torch.load(args.feat_file)

with open(args.cls_file, 'r') as f:
    cls_data = yaml.load(f, Loader=yaml.Loader)

feat_arr=torch.stack(list(text_feat.values()), dim=0).numpy()
cls_arr=np.array(list(cls_data.values()), dtype=int)

n_cls=20

print('start k-means')
k_means = KMeans(n_clusters=n_cls, random_state=42)
k_means.fit(feat_arr)
print(k_means.cluster_centers_)

pred = k_means.predict(feat_arr)

@numba.jit(nopython=True)
def get_pred_map(n_cls, pred, gt):
    pred_label_map=np.zeros((n_cls,2), dtype=np.float32) # gt -> pred
    for u in range(n_cls):
        rate_max = 0
        rate_idx = 0

        tmp_gt = np.zeros_like(gt)
        tmp_gt[gt == u] = 1

        for i in range(n_cls):
            tmp_pred=np.zeros_like(gt)
            pred_i = pred==i
            tmp_pred[pred_i] = 1
            rate = np.sum(tmp_pred & tmp_gt)/np.sum(tmp_gt)
            if rate>rate_max:
                rate_max=rate
                rate_idx=i
        pred_label_map[u,0]=rate_max
        pred_label_map[u,1]=rate_idx
    return pred_label_map


pred_label_map = get_pred_map(n_cls, pred, cls_arr-1)
print(pred_label_map)

feat_arr_w_center = np.vstack((feat_arr, k_means.cluster_centers_))

print('start tsne')
tsne = TSNE(n_components=2, learning_rate='auto', init='random')
res=tsne.fit_transform(feat_arr)  # 进行数据降维,并返回结果

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

colors = plt.cm.hsv(np.linspace(0, 1, 20)).tolist()

plt.figure(dpi=200)
for i in range(1,n_cls+1):
    d = res[cls_arr == i, :]  # 找出聚类类别为0的数据对应的降维结果
    plt.plot(d[:,0], d[:,1], 's', markerfacecolor='none', color=colors[i-1])
    #plt.plot(d[-n_cls+i-1,0], d[-n_cls+i-1,1], 'x', color=colors[i])

    d = res[pred == int(pred_label_map[i-1, 1]), :]
    plt.plot(d[:,0], d[:,1], 'x', color=colors[i-1])

plt.savefig("text_vis.png")
plt.show()